#include "Platform.hpp"

#include "cybertron/DefinesCore.hpp"
#include "cybertron/core/UtilString.hpp"
#include "cybertron/core/UtilFile.hpp"
#include "cybertron/core/UtilPath.hpp"

#include <unordered_set>

#ifdef CYBERTRON_WIN
#include <windows.h>
#endif

void* Platform::GetDllHandle(const std::string& filename)
{
	void* handle = nullptr;

#ifdef CYBERTRON_WIN
	if (!filename.empty()) {
		std::wstring make_wstring = cybertron::UtilString::u2w(filename);
		HINSTANCE hGetProcIDDLL = LoadLibrary(make_wstring.c_str());
		if (!hGetProcIDDLL) {
			DWORD err = GetLastError();
			printf("Error code: %d", err);
			// system error codes could be found here:
			// https://docs.microsoft.com/en-us/windows/win32/debug/system-error-codes
		}
		handle = (void*)hGetProcIDDLL;
	}
#endif

	return handle;
}

void Platform::FreeDllHandle(void* DllHandle)
{
#ifdef CYBERTRON_WIN
	// It is okay to call FreeLibrary on 0
	::FreeLibrary((HMODULE)DllHandle);
#endif
}

void* Platform::GetDllExport(void* DllHandle, const char* ProcName)
{
	void* proc = nullptr;
#ifdef CYBERTRON_WIN
	proc = (void*)::GetProcAddress((HMODULE)DllHandle, ProcName);
#endif

	return proc;
}

static const void *MapRvaToPointer(const IMAGE_DOS_HEADER *Header, const IMAGE_NT_HEADERS *NtHeader, size_t Rva)
{
	const IMAGE_SECTION_HEADER *SectionHeaders = (const IMAGE_SECTION_HEADER*)(NtHeader + 1);
	for (size_t SectionIdx = 0; SectionIdx < NtHeader->FileHeader.NumberOfSections; SectionIdx++)
	{
		const IMAGE_SECTION_HEADER *SectionHeader = SectionHeaders + SectionIdx;
		if (Rva >= SectionHeader->VirtualAddress && Rva < SectionHeader->VirtualAddress + SectionHeader->SizeOfRawData)
		{
			return (const BYTE*)Header + SectionHeader->PointerToRawData + (Rva - SectionHeader->VirtualAddress);
		}
	}
	return NULL;
}

static bool ReadLibraryImportsFromMemory(const IMAGE_DOS_HEADER* Header, std::vector<std::string>& ImportNames)
{
	bool bResult = false;
	if (Header->e_magic == IMAGE_DOS_SIGNATURE)
	{
		IMAGE_NT_HEADERS *NtHeader = (IMAGE_NT_HEADERS*)((BYTE*)Header + Header->e_lfanew);
		if (NtHeader->Signature == IMAGE_NT_SIGNATURE)
		{
			// Find the import directory header
			IMAGE_DATA_DIRECTORY *ImportDirectoryEntry = &NtHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

			// Enumerate the imports
			IMAGE_IMPORT_DESCRIPTOR *ImportDescriptors = (IMAGE_IMPORT_DESCRIPTOR*)MapRvaToPointer(Header, NtHeader, ImportDirectoryEntry->VirtualAddress);
			for (size_t ImportIdx = 0; ImportIdx * sizeof(IMAGE_IMPORT_DESCRIPTOR) < ImportDirectoryEntry->Size; ImportIdx++)
			{
				IMAGE_IMPORT_DESCRIPTOR *ImportDescriptor = ImportDescriptors + ImportIdx;
				if (ImportDescriptor->Name != 0)
				{
					const char *ImportName = (const char*)MapRvaToPointer(Header, NtHeader, ImportDescriptor->Name);
					ImportNames.push_back(ImportName);
				}
			}

			bResult = true;
		}
	}
	return bResult;
}

static bool ReadLibraryImports(const TCHAR* FileName, std::vector<std::string>& ImportNames)
{
	bool bResult = false;

	// Open the DLL using a file mapping, so we don't need to map any more than is necessary
	HANDLE NewFileHandle = CreateFile(FileName, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if (NewFileHandle != INVALID_HANDLE_VALUE)
	{
		HANDLE NewFileMappingHandle = CreateFileMapping(NewFileHandle, NULL, PAGE_READONLY, 0, 0, NULL);
		if (NewFileMappingHandle != NULL)
		{
			void* NewData = MapViewOfFile(NewFileMappingHandle, FILE_MAP_READ, 0, 0, 0);
			if (NewData != NULL)
			{
				const IMAGE_DOS_HEADER* Header = (const IMAGE_DOS_HEADER*)NewData;
				bResult = ReadLibraryImportsFromMemory(Header, ImportNames);
				UnmapViewOfFile(NewData);
			}
			CloseHandle(NewFileMappingHandle);
		}
		CloseHandle(NewFileHandle);
	}

	return bResult;
}

static std::string ConvertRelativePathToFull(const std::string& InPath)
{
	std::string fullPath;
	if (cybertron::UtilPath::isRelative(InPath)) {
		fullPath = cybertron::UtilPath::getWorkingPath() + InPath;
	}
	else {
		fullPath = InPath;
	}

	cybertron::UtilPath::normalizePath(fullPath);
	// need to collapse relative directoryies

	if (fullPath.size() == 0) {
		fullPath = "/";
	}

	return fullPath;
}

static bool ResolveImport(const std::string& Name, const std::vector<std::string>& SearchPaths, std::string& OutFileName)
{
	// Look for the named DLL on any of the search paths
	for (int Idx = 0; Idx < SearchPaths.size(); ++Idx)
	{
		std::string FileName = cybertron::UtilPath::combine(SearchPaths[Idx], Name);
		if (cybertron::UtilFile::exist(FileName))
		{
			OutFileName = ConvertRelativePathToFull(FileName);
			return true;
		}
	}
	return false;
}

static void ResolveMissingImportsRecursive(const std::string& FileName, const std::vector<std::string>& SearchPaths, std::vector<std::string>& ImportFileNames, std::unordered_set<std::string>& VisitedImportNames)
{
	// Read the imports for this library
	std::vector<std::string> ImportNames;
	std::wstring wFileName = cybertron::UtilString::u2w(FileName);
	if (ReadLibraryImports(wFileName.c_str(), ImportNames))
	{
		// Find all the imports that haven't already been resolved
		for (int Idx = 0; Idx < ImportNames.size(); Idx++)
		{
			const std::string& ImportName = ImportNames[Idx];
			if (!VisitedImportNames.count(ImportName))
			{
				// Prevent checking this import again
				VisitedImportNames.insert(ImportName);

				// Try to resolve this import
				std::wstring wImportName = cybertron::UtilString::u2w(ImportName);
				if (GetModuleHandle(wImportName.c_str()) == NULL)
				{
					std::string ImportFileName;
					if (ResolveImport(ImportName, SearchPaths, ImportFileName))
					{
						ResolveMissingImportsRecursive(ImportFileName, SearchPaths, ImportFileNames, VisitedImportNames);
						ImportFileNames.push_back(ImportFileName);
					}
				}
			}
		}
	}
}

static void LogImportDiagnostics(const std::string& FileName, const std::vector<std::string>& SearchPaths)
{
	std::vector<std::string> ImportNames;

	std::wstring wFileName = cybertron::UtilString::u2w(FileName);
	if (ReadLibraryImports(wFileName.c_str(), ImportNames))
	{
		bool bIncludeSearchPaths = false;
		for (const std::string& ImportName : ImportNames)
		{
			std::wstring wImportName = cybertron::UtilString::u2w(ImportName);
			if (GetModuleHandle(wImportName.c_str()) == nullptr)
			{
				printf("Missing import: %s\n", ImportName.c_str());
				bIncludeSearchPaths = true;
			}
		}
		if (bIncludeSearchPaths)
		{
			for (const std::string& searchPath : SearchPaths)
			{
				printf("Looked in: %s\n", searchPath.c_str());
			}
		}
	}
}

void* Platform::LoadLibraryWithSearchPaths(const std::string& FileName, const std::vector<std::string>& SearchPaths)
{
	// Make sure the initial module exists. If we can't find it from the path we're given, it's probably a system dll.
	std::string FullFileName = FileName;
	if (cybertron::UtilFile::exist(FullFileName))
	{
		// Convert it to a full path, since LoadLibrary will try to resolve it against the executable directory (which may not be the same as the working dir)
		FullFileName = ConvertRelativePathToFull(FullFileName);

		// Create a list of files which we've already checked for imports. Don't add the initial file to this list to improve the resolution of dependencies for direct circular dependencies of this
		// module; by allowing the module to be visited twice, any mutually depended on DLLs will be visited first.
		std::unordered_set<std::string> VisitedImportNames;

		// Find a list of all the DLLs that need to be loaded
		std::vector<std::string> ImportFileNames;
		ResolveMissingImportsRecursive(FullFileName, SearchPaths, ImportFileNames, VisitedImportNames);

		// Load all the missing dependencies first
		for (uint32_t Idx = 0; Idx < ImportFileNames.size(); ++Idx)
		{
			std::wstring fileName = cybertron::UtilString::u2w(ImportFileNames[Idx]);
			if (GetModuleHandle(fileName.c_str()) == nullptr)
			{
				if (LoadLibrary(fileName.c_str()))
				{
					//printf("Preloaded '%s'\n", ImportFileNames[Idx].c_str());
				}
				else
				{
					printf("Failed to preload '%s' (GetLastError=%d)", ImportFileNames[Idx].c_str(), GetLastError());
					LogImportDiagnostics(ImportFileNames[Idx], SearchPaths);
				}
			}
		}
	}

	// Try to load the actual library
	std::wstring wFullFileName = cybertron::UtilString::u2w(FullFileName);
	void* Handle = LoadLibrary(wFullFileName.c_str());
	if (Handle) {
		//printf("Loaded %s", FullFileName.c_str());
	}
	else {
		printf("Failed to load '%s' (GetLastError=%d)", FileName.c_str(), ::GetLastError());
		if (cybertron::UtilFile::exist(FileName)) {
			LogImportDiagnostics(FileName, SearchPaths);
		}
		else {
			printf("File '%s' does not exist", FileName.c_str());
		}
	}
	return Handle;
}
